import torch
import torch.nn as nn


class InstanceSeg(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.centroid_mlp = nn.Sequential(
            nn.Conv1d(128, 64, kernel_size=1, bias=False),
            nn.BatchNorm1d(64),
            nn.ReLU(True),
            nn.Conv1d(64, 3, kernel_size=1)
        )

        self.embedded_fc = nn.Conv1d(128, config.embedding_size)

    def forward(self, point_feature):
        centroid = self.centroid_mlp(point_feature)
        embedded = self.embedded_fc(point_feature)
        embedded = embedded.transpose(2, 1)
        return centroid, embedded





